from typing import List, Tuple, Dict, Any, Optional, Union
import os
import json
import logging
from PIL import Image
import torch
from torch.utils.data import Dataset
import torchvision.transforms.functional as TF  # 新增：图像张量操作

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 导入已有的process_image函数
from src.dataset import process_image

class VSTARBenchDataset(Dataset):
    def __init__(self, data_dir: str = "data/vstar_bench", 
                 subsets: List[str] = ["OCR", "direct_attributes", "relative_position", "GPT4V-hard"],
                 image_resolution: Optional[str] = "small",
                 max_samples: Optional[int] = None):
        """
        初始化VSTAR Bench数据集
        
        Args:
            data_dir: 数据集根目录
            subsets: 要加载的子集列表，如["OCR", "direct_attributes", "relative_position", "GPT4V-hard"]
            image_resolution: 图像分辨率 (high/mid/small/low)
            max_samples: 最大样本数，None表示加载所有样本
        """
        self.data_dir = data_dir
        self.subsets = subsets
        self.image_resolution = image_resolution
        self.max_samples = max_samples
        
        # 存储所有样本路径的列表
        self.sample_paths = []
        
        # 加载数据集样本路径
        self._load_dataset_paths()
        
        # 图像标记
        self.image_token = "<|image_1|>"
        
        logger.info(f"成功加载VSTAR Bench数据集，包含{len(self.sample_paths)}个样本")

    def _load_dataset_paths(self) -> None:
        """加载所有子集的JSON文件路径"""
        for subset in self.subsets:
            subset_dir = os.path.join(self.data_dir, subset)
            if not os.path.exists(subset_dir):
                logger.warning(f"子集目录不存在: {subset_dir}")
                continue
            
            logger.info(f"加载子集: {subset}")
            
            # 获取该子集中的所有JSON文件
            json_files = []
            for file in os.listdir(subset_dir):
                if file.endswith(".json"):
                    json_files.append(os.path.join(subset_dir, file))
            
            # 如果设置了最大样本数，限制样本数量
            if self.max_samples is not None and len(json_files) > self.max_samples:
                json_files = json_files[:self.max_samples]
                
            # 添加到总样本路径列表
            self.sample_paths.extend(json_files)
        
        if not self.sample_paths:
            raise ValueError("未能加载任何数据集样本")

    def _get_image(self, json_file_path: str) -> Optional[Image.Image]:
        """根据JSON文件路径获取对应的图像"""
        # 获取不带扩展名的文件名
        base_name = os.path.splitext(os.path.basename(json_file_path))[0]
        # 构建可能的图像路径
        subset_dir = os.path.dirname(json_file_path)
        
        # 尝试不同的图像扩展名
        image_extensions = [".jpg", ".jpeg", ".png", ".JPG", ".PNG", ".webp", ".WEBP"]
        for ext in image_extensions:
            image_path = os.path.join(subset_dir, f"{base_name}{ext}")
            if os.path.exists(image_path):
                try:
                    image = Image.open(image_path).convert("RGB")
                    return process_image(image, self.image_resolution), image_path
                except Exception as e:
                    logger.warning(f"加载图像{image_path}失败: {str(e)}")
                    return None, ""
        
        logger.warning(f"未找到与{json_file_path}对应的图像文件")
        return None, ""

    def __len__(self) -> int:
        return len(self.sample_paths)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        """获取单个样本"""
        json_file_path = self.sample_paths[idx]
        
        try:
            # 读取JSON文件
            with open(json_file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 提取字段并映射
            query_text = data.get("question", "")
            bbox = data.get("bbox", None)
            bbox = bbox[0] if bbox else None
            # target_object可能是列表，转换为字符串
            target_object = data.get("target_object", [])
            if isinstance(target_object, list):
                pos_text = ", ".join(target_object)
            else:
                pos_text = str(target_object)
            
            # 获取图像
            query_image, query_image_path = self._get_image(json_file_path)
            masked_image = None
            if bbox:
                masked_image = process_image(generate_masked_image(Image.open(query_image_path), bbox), self.image_resolution)
            # 确保至少有文本或图像
            if not query_text and not query_image and not bbox:
                logger.warning(f"样本{idx}包含空输入，跳过")
                # 返回下一个有效样本
                return self.__getitem__((idx + 1) % len(self))
            
            # 按照要求返回格式，其他字段设置为None
            return {
                "bbox": bbox,
                "masked_image": masked_image,
                "query_text": query_text,
                "query_image": query_image,
                "query_image_path": query_image_path,  # 存储相对路径
                "pos_text": pos_text,
                "pos_image": None,  # 其他字段设置为None
                "pos_image_path": "",
                "neg_text": "",
                "neg_image": None,
                "neg_image_path": ""
            }
            
        except Exception as e:
            logger.error(f"处理样本{json_file_path}时出错: {str(e)}")
            # 返回下一个有效样本
            return self.__getitem__((idx + 1) % len(self))

# 用于测试的函数
def inspect_dataset(dataset, num_samples=3):
    """查验数据集的几个样本"""
    for i in range(min(num_samples, len(dataset))):
        sample = dataset[i]
        print(f"\n样本 {i+1}:")
        print(f"查询文本: {sample['query_text']}")
        print(f"正文本: {sample['pos_text']}")
        print(f"查询图像: {'已加载' if sample['query_image'] is not None else '未加载'}")
        print(f"图像路径: {sample['query_image_path']}")

def generate_masked_image(original_image: Image.Image, bbox: Tuple[int, int, int, int], mask_color: Tuple[int, int, int] = (127, 127, 127)) -> Image.Image:
    """
    生成bbox区域被遮罩的图像
    Args:
        original_image: 原始PIL图像（query_image）
        bbox: (x, y, h, w) -> 左上角x坐标、左上角y坐标、bbox高度、bbox宽度
        mask_color: 遮罩颜色（默认灰色，RGB格式）
    Returns:
        masked_image: bbox区域被遮罩后的PIL图像
    """
    # print(bbox)
    x, y, h, w = bbox
    # 转换为PIL图像可操作的坐标（确保不超出图像边界）
    img_width, img_height = original_image.size
    x_end = min(x + w, img_width)  # 避免bbox右边界超出图像
    y_end = min(y + h, img_height)  # 避免bbox下边界超出图像
    x = max(x, 0)  # 避免x为负
    y = max(y, 0)  # 避免y为负

    # 将PIL图像转换为张量（便于区域替换）
    img_tensor = TF.to_tensor(original_image).unsqueeze(0)  # [1, 3, H, W]
    mask_tensor = TF.to_tensor(Image.new('RGB', (w, h), mask_color)).unsqueeze(0)  # [1, 3, h, w]

    # 遮罩bbox区域（替换对应张量切片）
    img_tensor[:, :, y:y_end, x:x_end] = mask_tensor[:, :, :y_end - y, :x_end - x]

    # 转换回PIL图像
    masked_image = TF.to_pil_image(img_tensor.squeeze(0))
    return masked_image

if __name__ == "__main__":
    # 示例用法
    try:
        # 创建数据集实例
        dataset = VSTARBenchDataset(
            data_dir="data/vstar_bench",
            subsets=["OCR", "direct_attributes"],  # 可选择要加载的子集
            image_resolution="small",
            max_samples=10  # 仅加载10个样本用于测试
        )
        
        # 查验数据集
        inspect_dataset(dataset, num_samples=3)
        
        # 打印数据集大小
        print(f"\n数据集总样本数: {len(dataset)}")
        
    except Exception as e:
        print(f"运行出错: {str(e)}")